#!/usr/bin/python
# -*- coding: utf-8 -*-
"""纯真ip数据库辅助工具

数据库格式说明: http://lumaqq.linuxsir.org/article/qqwry_format_detail.html
固定下载地址: http://update.cz88.net/soft/qqwry.rar
"""
import struct
import re
import os
from datetime import date
import urllib2
import shutil


class BufferFile(object):
    """为文件对象实现字符串的基本操作"""
    def __init__(self, filepath):
        self.ofile = open(filepath, 'rb')
        
    def __getslice__(self, i, j):
        self.ofile.seek(i, 0)
        s = self.ofile.read(j - i)
        return s
        
    def __del__(self):
        if self.ofile and hasattr(self.ofile, 'close'):
            self.ofile.close()
#            print 'ofile close'

class IPFormatError(Exception):
    """IP格式错误，目前只支持IPv4"""

class IPSearch(object):
    INDEX_START = (0, 4, '<L') # 第一个索引的开始位置
    INDEX_END = (4, 8, '<L') # 最后一个索引的开始位置
    INDEX = (7, '<L3B') # 索引数据结构
    MODE = (1, '<B') # 模式
    STRING_END = '<B' # 字符串结束标识位
    ADDRESS = (3, '<3B') # 3 bytes long
    IP = (4, '<L') # ip数据结构
    CHARSET = 'gbk' # 默认字符串编码
    REDIRECT_MODE_1 = 1 # 跳转模式1
    REDIRECT_MODE_2 = 2 # 跳转模式2
    CODE_PATH = os.path.dirname(os.path.abspath(__file__))
    DB_PATH = os.path.join(CODE_PATH, 'qqwry.dat')
    DOWN_URL = 'http://update.cz88.net/soft/qqwry.rar'
    
    def __init__(self, db_path=None, high_speed=False):
        """high_speed=True时，将所有索引一次性加载到内存中"""
        if db_path is not None:
            self.DB_PATH = db_path
        self._buf = None
        self._high_speed = high_speed
        if os.path.exists(self.DB_PATH):
            self.__load_db()
    
    def __load_db(self):
        if self._buf is not None:
            del self._buf
        if self._high_speed:
            f = open(self.DB_PATH, 'rb')
            try:
                self._buf = f.read()
            finally:
                f.close()
        else:
            self._buf = BufferFile(self.DB_PATH)
        self.index_start = struct.unpack(self.INDEX_START[2], 
                                    self._buf[self.INDEX_START[0]:self.INDEX_START[1]])[0]
        self.index_end = struct.unpack(self.INDEX_END[2], 
                                  self._buf[self.INDEX_END[0]:self.INDEX_END[1]])[0]
        # 加载ip数据库的版本信息
        self.total = (self.index_end - self.index_start) / self.INDEX[0] + 1
        self.version = self._get_location(self.index_end, True)[1]
        self.full_indexs = {}
        
    def update(self):
        """自动检测是否需要更新"""
        need_update = False
        if not os.path.exists(self.DB_PATH):
            need_update = True
        else:
            m = re.search(ur'(\d+)年(\d+)月(\d+)日', self.version)
            if m:
                y, m, d = int(m.group(1)), int(m.group(2)), int(m.group(3))
                last_date = date(y, m, d)
                if (date.today() - last_date).days > 5: # 5天更新一次
                    need_update = True
            else:
                need_update = True
        if need_update:
            self.__update()
        
    def __update(self):
        from extract7z import extract
        try:
            response = urllib2.urlopen(self.DOWN_URL)
            tmp_dat = os.path.join(os.path.dirname(self.DB_PATH), 'qqwry_tmp.rar')
            f = open(tmp_dat, 'wb')
            try:
                while True:
                    data = response.read(1024**2)
                    if not data:
                        break
                    f.write(data)
            finally:
                f.close()
            if self._buf is not None: # 先关闭io
                del self._buf
                self._buf = None
            extract_folder = extract(tmp_dat)
            dat_path = os.path.join(extract_folder, 'QQWry.Dat')
            if os.path.exists(dat_path):
                if os.path.exists(self.DB_PATH):
                    os.remove(self.DB_PATH)
                shutil.move(dat_path, self.DB_PATH)
            if os.path.exists(tmp_dat):
                os.remove(tmp_dat)
        except Exception, e:
            print e
        self.__load_db()
    
    def search(self, ip):
        """搜索ip的相关信息, 二分查找"""
        if self._high_speed and not self.full_indexs: # 加载索引
            self._load_full_index()
        ip_long = self._ip_to_long(ip)
        i, j = self.index_start, self.index_end
        pos = None
        while i < j:
            m = self._get_middle_offset(i, j)
            r = self._compare(ip_long, m)
            if r == 1:
                i = m + self.INDEX[0]
            elif r == -1:
                j = m - self.INDEX[0]
            else: # 已找到
                pos = m
                break
        # 此时, i == j
        if pos is None:
            r = self._compare(ip_long, i)
            if r == 0: # 找到
                pos = i
        if pos is not None:
            return self._get_location(pos)
        return None, None # 未找到
    
    def _load_full_index(self):
        """加载全索引到内存中"""
        self.full_indexs = {}
        index, end = self.index_start, self.index_end
        while index < end:
            self.full_indexs[index] = self._get_ip_range(index)
            index += self.INDEX[0]
            
    def _get_middle_offset(self, start, end):
        """获取中间位置"""
        c = (end - start) / self.INDEX[0] /2
        if c == 0:
            return end
        return start + c * self.INDEX[0]
        
    def _compare(self, ip_long, offset):
        """比较ip，判断ip是否在offset范围中
        1: 大于offset
        -1: 小于offset
        0: 在offset中
        """
        start_ip, end_ip = self._get_ip_range(offset)
        if ip_long < start_ip:
            return -1
        elif ip_long > end_ip:
            return 1
        else:
            return 0
    
    def _ip_to_long(self, ip):
        """获取ip的长整形格式"""
        ips = ip.split('.')
        if len(ips) < 4:
            raise IPFormatError(ip)
        l = int(ips[3])
        l |= (int(ips[2]) << 8) & 0xff00
        l |= (int(ips[1]) << 16) & 0xff0000
        l |= (int(ips[0]) << 24) & 0xff000000
        return l
    
    def _long_to_ip(self, l):
        return '%d.%d.%d.%d' % ((l & 0xff000000) >> 24,
                                (l & 0x00ff0000) >> 16,
                                (l & 0x0000ff00) >> 8,
                                l & 0x000000ff,
                                )
        
    def _get_ip_range(self, offset):
        """获取起始ip和结束ip，长整形格式"""
        if offset in self.full_indexs:
            return self.full_indexs[offset]
        start_ip, data_offset = self._get_index(offset)
        end_ip = struct.unpack('<L', self._buf[data_offset:data_offset+self.IP[0]])[0]
        return start_ip, end_ip
    
    def _get_index(self, offset):
        """获取索引"""
        buf = self._buf[offset:offset+self.INDEX[0]]
        index = struct.unpack(self.INDEX[1], buf)
        start_ip, data_offset = index[0], self._get_3byte_long(index[1:])
        return start_ip, data_offset
        
    def _get_location(self, offset, deal_mode_0=False):
        """获取区域信息, 包含国家信息和区域信息"""
        data_offset = self._get_index(offset)[1]
        offset = data_offset + self.IP[0]
        end = offset + self.MODE[0]
        mode = struct.unpack(self.MODE[1], self._buf[offset:end])[0]
        if mode == self.REDIRECT_MODE_1: # 模式1
            offset = end
            end = offset + self.ADDRESS[0]
            address = struct.unpack(self.ADDRESS[1], self._buf[offset:end])
            offset = self._get_3byte_long(address)
            end = offset+self.MODE[0]
            mode = struct.unpack(self.MODE[1], self._buf[offset:end])[0]
            if mode == self.REDIRECT_MODE_2: # 模式2
                offset = end
                end = offset+self.ADDRESS[0]
                address = struct.unpack(self.ADDRESS[1], self._buf[offset:end])
                country = self._get_string(self._get_3byte_long(address)) # 读取国家信息
                area_offset = offset + self.ADDRESS[0]
            else:
                end = self._get_string_end(offset)
                country = self._get_string(offset, end)
                area_offset = end
        elif mode == self.REDIRECT_MODE_2: # 模式2
            offset = end
            end = offset + self.ADDRESS[0]
            address = struct.unpack(self.ADDRESS[1], self._buf[offset:end])
            country = self._get_string(self._get_3byte_long(address))
            area_offset = offset + self.ADDRESS[0]
        else: # 普通模式
            end = self._get_string_end(offset)
            country = self._get_string(offset, end)
            area_offset = end
        area = self._get_area(area_offset, deal_mode_0)
        if not area:
            area = None
        return country, area
    
    def _get_area(self, offset, deal_mode_0=False):
        """获取区域信息"""
        mode = struct.unpack(self.MODE[1], self._buf[offset:offset+self.MODE[0]])[0]
        if mode in (self.REDIRECT_MODE_1, self.REDIRECT_MODE_2):
            offset += self.MODE[0]
            address = struct.unpack(self.ADDRESS[1], self._buf[offset:offset+self.ADDRESS[0]])
            offset = self._get_3byte_long(address)
            if offset == 0:
                return None
        elif mode == 0 and deal_mode_0:
            offset += self.MODE[0]
        end = self._get_string_end(offset)
        area = self._get_string(offset, end)
        return area

    def _get_string_end(self, offset):
        """获取字符串的结束位置"""
        end = offset + 1
        while True:
            i = struct.unpack(self.STRING_END, self._buf[offset:end])[0]
            if i == 0:
                break
            offset += 1
            end = offset + 1
        return end - 1

    def _get_string(self, offset, end=None):
        """获取字符串"""
        if end is None:
            end = self._get_string_end(offset)
        s = self._buf[offset:end]
        return s.decode(self.CHARSET)
        
    def _get_3byte_long(self, chars):
        """获取由3个字节组成的长整形"""
        l = chars[0]
        l |= (chars[1] << 8) & 0xff00
        l |= (chars[2] << 16) & 0xff0000
        return l
    
    RE_PROVINCE_CITY = re.compile(ur'(.+?省|(?:上海|北京|天津|重庆)市)(?:(.+?)市)?')
    RE_CITY = re.compile(u'(.+?)市', re.IGNORECASE)
    RE_EDU = re.compile(ur'(.+?(?:大学|学院))|(.+?中学)|(.+?小学)|(.+?学校)')
    RE_INTERNETCAFE = re.compile(ur'(.+?网吧)')
    RE_LIBRARY = re.compile(ur'(.*?图书馆)')
    RE_GOV = re.compile(ur'(.*?办公室|.*?政府|.*?公安局)')
    RE_COMPANY = re.compile(ur'(.*?公司)')
    
    RE_REPLACE = re.compile(ur'\(.+\)')
    
    PROVINCES_2 = (
                 u'北京', u'上海', u'天津', u'重庆' # 直辖市
                 u'河北', u'河南',
                 u'湖北', u'湖南',
                 u'江苏', u'江西',
                 u'辽宁', u'吉林',
                 u'陕西',
                 u'山西', u'山东',
                 u'四川',
                 u'青海',
                 u'安徽',
                 u'广东', u'浙江', u'福建', # 沿海省份
                 u'台湾', u'海南', # 岛
                 u'贵州', u'甘肃', u'云南',
                 u'西藏', u'宁夏', u'广西', u'新疆', # 自治区
                 u'香港', u'澳门', # 特别行政区
                 )
    PROVINCES_3 = (u'黑龙江', u'内蒙古',)
    def translate(self, country, area):
        """将获取到的国家和地区信息，翻译成可用的信息
        {'province': p, # 省, 直辖市
         'city': city, # 市
         'university': school, # 学校，大学
         'middle_school': school, # 中学
         'primary_school': primary_school, # 小学
         'other_school': other_school, # 其他学校
         'internet_cafe': internetcafe, # 网吧
         'library': library, # 图书馆
         'gov': gov, # 政府, 办公室, 公安局
         'company': company, # 公司
         }
        """
        result = {}
        # province and city
        province, city = None, None
        m = self.RE_PROVINCE_CITY.search(country)
        if m:
            province, city = m.groups()
            if province:
                province = province[:-1]
        elif country[:2] in self.PROVINCES_2:
            province = country[:2]
            m = self.RE_CITY.search(country[2:])
            if m:
                city = m.group(1)
        elif country[:3] in self.PROVINCES_3:
            province = country[:3]
            m = self.RE_CITY.search(country[3:])
            if m:
                city = m.group(1)
        if province:
            if u'淅江' in province:
                province = province.replace(u'淅江', u'浙江')
            result['province'] = province
        if city:
            result['city'] = city
        # 网吧
        internet_cafe = None
        m = None
        if area:
            m = self.RE_INTERNETCAFE.search(area)
            if m:
                internet_cafe = u'%s%s' % (country, m.group(1))
        if not m:
            m = self.RE_INTERNETCAFE.search(country)
            if m:
                internet_cafe = m.group(1)
        if internet_cafe:
            result['internet_cafe'] = self.RE_REPLACE.sub(u'', internet_cafe)
        if not internet_cafe: # 没发现时网吧，才去找学校
        # 学校
            university, middle_school, primary_school, other_school = None, None, None, None
            m = None
            if area:
                m = self.RE_EDU.search(area)
                if m:
                    university, middle_school, primary_school, other_school = m.groups()
            if not m:
                m = self.RE_EDU.search(country)
                if m:
                    university, middle_school, primary_school, other_school = m.groups()
            if university:
                result['university'] = university
            if middle_school:
                result['middle_school'] = u'%s%s' % (country, middle_school)
            if primary_school:
                result['primary_school'] = u'%s%s' % (country, primary_school)
            if other_school:
                result['other_school'] = u'%s%s' % (country, other_school)
            # 图书馆
            library = None
            m = None
            if area:
                m = self.RE_LIBRARY.search(area)
                if m:
                    library = u'%s%s' % (country, m.group(1))
            if not m:
                m = self.RE_LIBRARY.search(country)
                if m:
                    library = m.group(1)
            if library:
                result['library'] = library
                
            gov = None
            m = None
            if area:
                m = self.RE_GOV.search(area)
                if m:
                    gov = u'%s%s' % (country, m.group(1))
            if not m:
                m = self.RE_GOV.search(country)
                if m:
                    gov = m.group(1)
            if gov:
                result['gov'] = gov
                
            company = None
            m = None
            if area:
                m = self.RE_COMPANY.search(area)
                if m:
                    company = m.group(1)
            if not m:
                m = self.RE_GOV.search(country)
                if m:
                    company = m.group(1)
            if company:
                result['company'] = company
        return result
    
    
def test():
    ipsearch = IPSearch(high_speed=False)
    ipsearch.update()
    print ipsearch.total, ipsearch.version
    test_ips = [
                '0.0.0.0', '24.235.129.1', '219.150.200.75',
                '121.11.14.153',
                '60.28.249.179', '60.28.249.167', '60.28.249.168',
                '172.254.16.5', 
                '127.0.0.1', '127.0.0.2', '127.1.0.1', '127.0.2.1', 
                '10.20.208.60', '10.20.208.62', '192.168.47.1',
                '255.255.255.255', # 版本信息
                '225.0.0.0',
                '116.230.255.252', '202.120.57.221', # 上海
                '222.203.146.1', '58.19.238.89', # 小学
                ]
    for ip in test_ips:
        print ip.center(50, '*')
        c, a = ipsearch.search(ip)
        print c, '---', a
        print u'%s: %s | %s' % (ip, c, a)
        r = ipsearch.translate(c, a)
        for k, v in r.iteritems():
            print '%s: %s' % (k, v.encode('utf-8'))
    try:
        ipsearch.search('12.23')
    except IPFormatError:
        pass
    
    
if __name__ == '__main__':
    test()